import os
import zipfile
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


def load_metr_la_data():
    """
    加载原始数据，并归一化后返回
    """

    A = np.load("./data/B.npy")  # A=float32 x=float64
    B = np.load("./data/all_data_0224.npz", allow_pickle=True)
    X = B['arr_0'].transpose((1, 2, 0))
    X = X.astype(np.float32) #X =（303,4,17520）
    X = X[:, :, :3360]

    # Normalization using Z-score method
    means = np.mean(X, axis=(0, 2))
    X = X - means.reshape(1, -1, 1)
    stds = np.std(X, axis=(0, 2))
    X = X / stds.reshape(1, -1, 1)

    return A, X, means, stds


def get_normalized_adj(A):
    """
    Returns the degree normalized adjacency matrix.
    """
    A = A + np.diag(np.ones(A.shape[0], dtype=np.float32))
    D = np.array(np.sum(A, axis=1)).reshape((-1,))
    D[D <= 10e-5] = 10e-5    # Prevent infs
    diag = np.reciprocal(np.sqrt(D))
    A_wave = np.multiply(np.multiply(diag.reshape((-1, 1)), A),
                         diag.reshape((1, -1)))
    return A_wave



def i_delect (X,num_timesteps_input, num_timesteps_output):
    """
    制定sample在时间轴的切片
    """
    n=int(X.shape[2]/48)
    print(n)
    all_time_stamps = list(range(int(X.shape[2])))  # 所有的时间戳
    all_start_time_stamps = [-1 + _ * 48 for _ in range(n)]     # 所有每天开始的时间戳
    h = []
    for start_time_stamp_now in all_time_stamps:  # start_time_stamp_now是当前读取的时间片的起点
        end_time_stamp_now = start_time_stamp_now + num_timesteps_input+num_timesteps_output # end_time_stamp_now是当前读取到的时间片的终点
        time_slot = list(range(start_time_stamp_now, end_time_stamp_now))  # 当前提取出的时间片段
        #print(time_slot)
        ignore_flag = False  # 是否忽略当前的时间偏短

    # 假如当前时间偏短内存在每日的开始的时间戳，该时间片段是需要被忽略的
        for start_time_stamps in all_start_time_stamps:
            if start_time_stamps in time_slot and start_time_stamp_now != start_time_stamps:    # 规避0 in [0, 37]、38 in [38, 75]...这种情况
                ignore_flag = True
                break
            else:
                pass

        if ignore_flag:
            continue
        else:
        # 去提取数据并送入模型
            h.append(start_time_stamp_now)
            pass
    return h


def generate_dataset(X, num_timesteps_input, num_timesteps_output):
    """
    :param X: Node features of shape (num_timesteps, num_features,num_vertices)
    :return:
        - Node features divided into multiple samples. Shape is
          (num_samples, num_vertices, num_features, num_timesteps_input).
        - Node targets for the samples. Shape is
          (num_samples, num_vertices, num_features, num_timesteps_output).
    """

    i_data= i_delect(X,num_timesteps_input, num_timesteps_output)
    indices = [(i, i + (num_timesteps_input + num_timesteps_output)) for i
               in i_data]
    print(indices)
    # 根据之前制作的切片划分数据集
    features, target = [], []
    for i, j in indices:
        if j < X.shape[2]:
             features.append(X[:, :, i: i + num_timesteps_input].transpose(
                     (0, 2, 1)))
             target.append(X[:, 0, i + num_timesteps_input: j])
        else:
         pass
    return   torch.from_numpy(np.array(features)), \
             torch.from_numpy(np.array(target))
#最终得到      features.shape=(2379,303,12,4),
#             target.shape=(2379,303,3).
